#include <linux/module.h>
#include <linux/proc_fs.h>
#include <linux/slab.h>
#include <linux/uaccess.h>
#include <linux/usb.h>
#include <linux/kernel.h>

#include "usb_ctrl.h"

#define SEC_DEVCTL_DIR "sec_devctl"
#define SEC_USB_POLICY "usb_policy"
#define SEC_DEVCTL_SWITCH "devctl_switch"
struct proc_dir_entry *sec_devctl_dir;
struct proc_dir_entry *udev_policy_file;
struct proc_dir_entry *dev_switch_file;

//interface对应子类的开关，默认功能关闭
static struct type_switch intf_type_switch[] = {
	{INTF_SERIAL,			0},
	{INTF_PARALLEL,			0},
	{INTF_1394,			0},
	{INTF_PCMCIA,			0},
	{INTF_TYPEC,			0},
	{INTF_HDMI,			0},
	{INTF_DISPLAYPORT,		0},
	{-1,				0}
};

//device对应子类的开关，默认功能开启
static struct type_switch device_type_switch[] = {
	{DEVICE_BLUETOOTH,			0},
	{DEVICE_INFRARED,			0},
	{DEVICE_MODEM,			0},
	{DEVICE_STORAGE_DRIVE,		0},
	{DEVICE_SCSI_RAID,			0},
	{DEVICE_WINCE,			0},
	{DEVICE_FLOPPY_DRIVE,		0},
	{DEVICE_IMAGE,			0},
	{DEVICE_MULTI_SERIAL_CARD,		0},
	{DEVICE_ETHERNET,			0},
	{DEVICE_WIRELESS,			0},
	{DEVICE_PRINTER,			0},
	{DEVICE_PRINTER_LOCAL,		0},
	{DEVICE_CDROM,			0},
	{DEVICE_VOLUMESHADOW,		0},
	{DEVICE_TAPE,			0},
	{-1,				0}
};

//usb对应子类的开关，默认功能开启
static struct type_switch usb_type_switch[] = {
        {USB_CLASS_PER_INTERFACE,       0},
        {USB_CLASS_AUDIO,               0},
        {USB_CLASS_COMM,                0},
        {USB_CLASS_HID,                 0},
        {USB_CLASS_PHYSICAL,            0},
        {USB_CLASS_STILL_IMAGE,         0},
        {USB_CLASS_PRINTER,             0},
        {USB_CLASS_MASS_STORAGE,        0},
        {USB_CLASS_HUB,                 0},
        {USB_CLASS_CDC_DATA,            0},
        {USB_CLASS_CSCID,               0},
        {USB_CLASS_CONTENT_SEC,         0},
        {USB_CLASS_VIDEO,               0},
        {USB_CLASS_WIRELESS_CONTROLLER, 0},
	{USB_CLASS_PERSONAL_HEALTHCARE,	0},
	{USB_CLASS_AUDIO_VIDEO,		0},
	{USB_CLASS_BILLBOARD,		0},
	{USB_CLASS_USB_TYPE_C_BRIDGE,	0},
        {USB_CLASS_MISC,                0},
        {USB_CLASS_APP_SPEC,            0},
        {USB_CLASS_VENDOR_SPEC,         0},
	{USB_CLASS_VENDOR_SPEC,		0},
	{-1,				0},
};

//usb白名单链表
static LIST_HEAD(usb_wlist);
//usb白名单链表读写所
DEFINE_RWLOCK(usb_wlist_rwlock);

//从白名单找到传入参数attr的数据
//返回1找到，0未找到
static int wlist_find(struct usb_attr *attr){
	struct list_head *pos = NULL;
	struct usb_wlist_node *info;
	int ret = 0;

	read_lock(&usb_wlist_rwlock);
	list_for_each(pos, &usb_wlist){
		info = list_entry(pos, struct usb_wlist_node, list);
		if(info->vid == attr->vid &&
				info->pid == attr->pid &&
				info->clsid == attr->clsid){
			ret = 1;
			goto output;
		}
	}

output:
	read_unlock(&usb_wlist_rwlock);
	return ret;
}


//遍历所有usb设备的interface找到并释放指定的data的interface驱动
static int foreach_udev_release_driver(struct usb_device *udev, void *data){
	struct usb_interface *intf;
	struct usb_attr *attr;
	int i;

	attr = (struct usb_attr *)data;
	if(attr->vid != udev->descriptor.idVendor || attr->pid != udev->descriptor.idProduct)
		return 0;

	for(i = 0; i < udev->actconfig->desc.bNumInterfaces; i++){
		intf = udev->actconfig->interface[i];
		if(intf){
			if(attr->clsid == intf->cur_altsetting->desc.bInterfaceClass &&
					intf->dev.driver)
				usb_driver_release_interface(to_usb_driver(intf->dev.driver), intf);
		}
	}

	return 0;
}

//释放attr对应的interface驱动
static int find_usb_release_driver(struct usb_attr *attr){
	usb_for_each_dev(attr, foreach_udev_release_driver);

	return 0;
}

//设置switch_data中类型对应的开关
static int type_switch_set(struct dev_switch *switch_data){
	struct type_switch *arr;
        int ix;
	int flag = 0;

	if(!switch_data)
		return -EFAULT;

	switch(switch_data->type){
		case DEV_TYPE_USB:
			arr = usb_type_switch;
			break;
		case DEV_TYPE_DEVICE:
			arr = device_type_switch;
			break;
		case DEV_TYPE_INTF:
			arr = intf_type_switch;
			break;
		default:
			return -EFAULT;
	}

        for (ix = 0; arr[ix].desc != -1; ix++)
               	if (arr[ix].desc == switch_data->sub_type){
			arr[ix].open = switch_data->open;
			flag = 1;
		}

	if(flag)
		return 0;
	else
		return -EFAULT;
}

//获取switch_data中指定类型的开关的值
static int type_switch_get(struct dev_switch *switch_data){
	struct type_switch *arr;
	int ix;

	if(!switch_data)
		return -EFAULT;

	switch(switch_data->type){
		case DEV_TYPE_USB:
			arr = usb_type_switch;
			break;
		case DEV_TYPE_DEVICE:
			arr = device_type_switch;
			break;
		case DEV_TYPE_INTF:
			arr = intf_type_switch;
			break;
		default:
			return -EFAULT;
	}

	for (ix = 0; arr[ix].desc != -1; ix++)
		if(arr[ix].desc == switch_data->sub_type)
			break;
	switch_data->open = arr[ix].open;
	return arr[ix].open;
}

//当有usb设备插入时，触发此函数，检测插入的usb设备的存储类中不在白名单的设备驱动卸载，只放行白名单设备
static int notify_dev_add_func(struct usb_device *udev){
	int i;
	struct usb_interface *intf;
	struct usb_attr attr = { 0 };
	int find = 0;
	struct dev_switch tmp = { 0 };

	attr.vid = udev->descriptor.idVendor;
	attr.pid = udev->descriptor.idProduct;
	for(i = 0; i < udev->actconfig->desc.bNumInterfaces; i++){
		intf = udev->actconfig->interface[i];
		if(intf){
			attr.clsid = intf->cur_altsetting->desc.bInterfaceClass;
			tmp.type = DEV_TYPE_USB;
			tmp.sub_type = attr.clsid;
			type_switch_get(&tmp);
			if(!tmp.open)
				return 0;

			switch(attr.clsid){
				case USB_CLASS_MASS_STORAGE:
					find = wlist_find(&attr);
					if(find)
						printk("zds detected vid=%04x pid=%04x clsid=%04x in white list: pass\n",
								attr.vid,
								attr.pid,
								attr.clsid);
					else
						printk("zds detected vid=%04x pid=%04x clsid=%04x in white list: stop\n",
								attr.vid,
								attr.pid,
								attr.clsid);
					if(!find && intf->dev.driver)
						usb_driver_release_interface(to_usb_driver(intf->dev.driver), intf);
					break;
				default:
					break;
			}
		}
	}

	return 0;
}

//设置usb notify
static int usb_notify(struct notifier_block *self, unsigned long action, void *dev){
	switch(action){
		case USB_DEVICE_ADD:
			notify_dev_add_func(dev);
			break;
		case USB_DEVICE_REMOVE:
			break;
	}
	return NOTIFY_OK;
}

static struct notifier_block usb_notify_block = {
	.notifier_call = usb_notify,
};


//data中struct all_usb_data为数据头，将查询结果放在数据体部分
static int all_usb_foreach(struct usb_device *udev, void *data){
	struct usb_interface *intf;
	struct usb_attr *attr_item;
	struct all_usb_data *all_usb;
	int i;

	if(!data)
		return 0;

	all_usb = (struct all_usb_data *)data;
	if(all_usb->index >= all_usb->data_len)
		return -ENOMEM;
		
	for(i = 0; i < udev->actconfig->desc.bNumInterfaces; i++){
		intf = udev->actconfig->interface[i];
		if(intf){
			if(all_usb->index + sizeof(struct usb_attr) > all_usb->data_len)
				return -ENOMEM;
			if(intf->cur_altsetting->desc.bInterfaceClass == USB_CLASS_MASS_STORAGE){
				printk("zds get all USB_CLASS_MASS_STORAGE:vid=%04x pid=%04x clsid=%04x\n",
						udev->descriptor.idVendor,
						udev->descriptor.idProduct,
						intf->cur_altsetting->desc.bInterfaceClass);
				attr_item = (struct usb_attr *)(all_usb->data + all_usb->index);	
				attr_item->vid = udev->descriptor.idVendor;
				attr_item->pid = udev->descriptor.idProduct;
				attr_item->clsid = intf->cur_altsetting->desc.bInterfaceClass;
				all_usb->index += sizeof(struct usb_attr);
			}
		}
	}


	return 0;
}

//读取所有usb数据和usb白名单数据
static ssize_t udev_ctrl_read(struct file *file, char __user *ubuf, size_t size, loff_t *off){
	int index = 0;
	char *buf;
	struct dev_data_hdr *hdr;
	struct list_head *pos = NULL;
	struct usb_wlist_node *info;
	struct usb_attr *attr_item;
	struct all_usb_data *all_usb;
	int ret = 0;
	int wlist_count = 0;

	buf = kzalloc(PAGE_SIZE, GFP_KERNEL);
	if(!buf)
		return -ENOMEM;

	hdr = (struct dev_data_hdr *)buf;
	*(int *)(buf + sizeof(struct dev_data_hdr)) = PAGE_SIZE - sizeof(struct dev_data_hdr);
	index += sizeof(struct dev_data_hdr);
	all_usb = (struct all_usb_data *)(buf + index);
	all_usb->data_len = PAGE_SIZE - sizeof(struct dev_data_hdr) - sizeof(struct all_usb_data);
	all_usb->index = 0;
	usb_for_each_dev(all_usb, all_usb_foreach);
	hdr->usb_nu = all_usb->index / sizeof(struct usb_attr);
	memcpy(buf + index, all_usb->data, hdr->usb_nu * sizeof(struct usb_attr));
	index += hdr->usb_nu * sizeof(struct usb_attr);

	read_lock(&usb_wlist_rwlock);
	list_for_each(pos, &usb_wlist){
		info = list_entry(pos, struct usb_wlist_node, list);
		
		attr_item = (struct usb_attr *)(buf + index);
		attr_item->vid = info->vid;
		attr_item->pid = info->pid;
		attr_item->clsid = info->clsid;
		printk("zds get white list USB_CLASS_MASS_STORAGE:vid=%04x pid=%04x clsid=%04x\n",
				attr_item->vid,
				attr_item->pid,
				attr_item->clsid);
		index += sizeof(struct usb_attr);
		wlist_count++;

		if(index + sizeof(struct usb_attr) > PAGE_SIZE){
			ret = -ENOMEM;
			goto output;
		}
	}
	read_unlock(&usb_wlist_rwlock);
	hdr->usb_wlist_nu = wlist_count;


	if(0 != copy_to_user(ubuf, buf, index)){
		ret = -EFAULT;
		goto output;
	}

	ret = sizeof(struct dev_data_hdr) + 
		hdr->usb_nu * sizeof(struct usb_attr) +
		hdr->usb_wlist_nu * sizeof(struct usb_attr);
output:
	kfree(buf);
	return ret;
}

//将attr的usb数据添加到白名单中
static int usb_add_wlist(struct usb_attr *attr){
	struct usb_wlist_node *info, *new;
	struct list_head *pos = NULL;

	printk("zds add vid=%04x pid=%04x clsid=%04x to white list\n", attr->vid, attr->pid, attr->clsid);
	new = kzalloc(sizeof(struct usb_wlist_node), GFP_KERNEL);
	if(!new){
		return -ENOMEM;
	}
	new->vid = attr->vid;
	new->pid = attr->pid;
	new->clsid = attr->clsid;

	write_lock(&usb_wlist_rwlock);
	list_for_each(pos, &usb_wlist){
		info = list_entry(pos, struct usb_wlist_node, list);
		if(info->vid == attr->vid &&
				info->pid == attr->pid &&
				info->clsid == attr->clsid){
			kfree(new);
			write_unlock(&usb_wlist_rwlock);
			return 0;
		}
	}	
	list_add(&new->list, &usb_wlist);
	write_unlock(&usb_wlist_rwlock);

	return 0;
}

//将attr数据从白名单删除
static int usb_del_wlist(struct usb_attr *attr){
	struct usb_wlist_node *item = NULL, *tmp = NULL;

	printk("zds remove vid=%04x pid=%04x clsid=%04x to white list\n", attr->vid, attr->pid, attr->clsid);
	write_lock(&usb_wlist_rwlock);
	list_for_each_entry_safe(item, tmp, &usb_wlist, list){
		if(attr->vid == item->vid &&
				attr->pid == item->pid &&
				attr->clsid == item->clsid){
			list_del(&item->list);
			kfree(item);
			break;
		}
	}
	write_unlock(&usb_wlist_rwlock);

	return 0;
}

//从procfs下发usb配置信息
static ssize_t udev_ctrl_write(struct file *file, const char __user *ubuf, size_t count, loff_t *off){
	char *tmp;
	struct usb_attr *attr;
	int err = 0, ret = 0;

	tmp = kzalloc(count + 1, GFP_KERNEL);
	if(!tmp)
		return -ENOMEM;

	if(0 != copy_from_user(tmp, ubuf, count))
		return -EFAULT;

	attr = (struct usb_attr *)(tmp + 1);
	switch(tmp[0]){
		case '+':
			err = usb_add_wlist(attr);
			break;
		case '-':
			err = usb_del_wlist(attr);
			if(err == 0)
				find_usb_release_driver(attr);
			break;
		default:
			ret = -EFAULT;
			break;
	}

	if(err == 0)
		ret = count;

	return ret;
}

const struct file_operations udev_ctrl_opt = {
	.read = udev_ctrl_read,
	.write = udev_ctrl_write,
};

static int udev_procfs_init(void){
	udev_policy_file = proc_create(SEC_USB_POLICY, 0666, sec_devctl_dir, &udev_ctrl_opt);
	if(!udev_policy_file){
		printk("can not create /proc/sec_dev_conf\n");
		return -EFAULT;
	}

	return 0;
}

static void udev_procfs_exit(void){
	if(udev_policy_file)
		proc_remove(udev_policy_file);
}

//读取开关配置
static ssize_t switch_read(struct file *file, char __user *ubuf, size_t size, loff_t *off){
	char *buf;
	int ret;
	struct dev_switch *item;
	int i;
	int count = 0;

	buf = kzalloc(PAGE_SIZE, GFP_KERNEL);
	if(!buf)
		return -ENOMEM;
	item = (struct dev_switch *)buf;

	for (i = 0; intf_type_switch[i].desc != -1; i++){
		(item + count)->type = DEV_TYPE_INTF;
		(item + count)->sub_type = intf_type_switch[i].desc;
		(item + count)->open = intf_type_switch[i].open;
		count++;
	}

	for(i = 0; device_type_switch[i].desc != -1; i++){
		(item + count)->type = DEV_TYPE_DEVICE;
		(item + count)->sub_type = device_type_switch[i].desc;
		(item + count)->open = device_type_switch[i].open;
		count++;
	}

	for(i = 0; usb_type_switch[i].desc != -1; i++){
		(item + count)->type = DEV_TYPE_USB;
		(item + count)->sub_type = usb_type_switch[i].desc;
		(item + count)->open = usb_type_switch[i].open;
		count++;
	}

	ret = count *sizeof(struct dev_switch);
	if(size < count * sizeof(struct dev_switch)){
		ret = -EFAULT;
		goto output;
	}

	if(0 != copy_to_user(ubuf, buf, count * sizeof(struct dev_switch))){
		ret = -EFAULT;
		goto output;
	}

output:
	kfree(buf);
	return ret;
}

//下发开关配置
static ssize_t switch_write(struct file *file, const char __user *ubuf, size_t count, loff_t *off){
	char *tmp;
	int ret;
	struct dev_switch *data;

	if(count < sizeof(struct dev_switch))
		return -EFAULT;

	tmp = kzalloc(count + 1, GFP_KERNEL);
	if(!tmp)
		return -ENOMEM;

	if(0 != copy_from_user(tmp, ubuf, count)){
		kfree(tmp);
		return -EFAULT;
	}

	data = (struct dev_switch *)tmp;
	ret = type_switch_set(data);
	if(!ret){
		kfree(tmp);
		return ret;
	}

	kfree(tmp);
	return count;
}

const struct file_operations sec_devctl_switch_opt = {
	.read = switch_read,
	.write = switch_write,
};

static int dev_switch_procfs_init(void){
	dev_switch_file = proc_create(SEC_DEVCTL_SWITCH, 0666, sec_devctl_dir, &sec_devctl_switch_opt);
	if(!dev_switch_file){
		printk("can not create file switch\n");
		return -EFAULT;
	}
		
	return 0;
}

static int __init dev_ctrl_init(void){
	int ret = 0;
	
	sec_devctl_dir = proc_mkdir(SEC_DEVCTL_DIR, NULL);
	if(!sec_devctl_dir)
		return -EFAULT;

	ret = udev_procfs_init();
	if(ret)
		goto rm_devctl_dir;

	ret = dev_switch_procfs_init();
	if(ret)
		goto rm_udev_policy;
	
	usb_register_notify(&usb_notify_block);

	return 0;
rm_udev_policy:
	proc_remove(udev_policy_file);
rm_devctl_dir:
	proc_remove(sec_devctl_dir);
	return ret;
}

static void dev_switch_procfs_exit(void){
	if(dev_switch_file)
		proc_remove(dev_switch_file);
}

static void __exit dev_ctrl_exit(void){
	struct usb_wlist_node *item = NULL, *tmp = NULL;

	usb_unregister_notify(&usb_notify_block);
	udev_procfs_exit();
	dev_switch_procfs_exit();
	if(sec_devctl_dir)
		proc_remove(sec_devctl_dir);

	write_lock(&usb_wlist_rwlock);
	list_for_each_entry_safe(item, tmp, &usb_wlist, list){
		list_del(&item->list);
		kfree(item);
	}
	write_unlock(&usb_wlist_rwlock);

}

module_init(dev_ctrl_init);
module_exit(dev_ctrl_exit);
MODULE_AUTHOR("zds");
MODULE_LICENSE("GPL");
